package cn.walket.test.db;

import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.ast.statement.SQLAlterTableAddIndex;
import com.alibaba.druid.sql.ast.statement.SQLAlterTableDropIndex;
import com.alibaba.druid.sql.ast.statement.SQLAlterTableItem;
import com.alibaba.druid.sql.ast.statement.SQLAlterTableRenameIndex;
import com.alibaba.druid.sql.ast.statement.SQLAlterTableStatement;
import com.alibaba.druid.sql.ast.statement.SQLTableElement;
import com.alibaba.druid.sql.dialect.mysql.ast.MySqlKey;
import com.alibaba.druid.sql.dialect.mysql.ast.MySqlPrimaryKey;
import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlCreateTableStatement;
import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlTableIndex;
import com.alibaba.druid.sql.dialect.mysql.parser.MySqlStatementParser;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.file.Files;
import java.util.ArrayList;
import java.util.List;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import lombok.Getter;
import org.springframework.core.io.DefaultResourceLoader;
import org.springframework.core.io.FileSystemResource;
import org.springframework.core.io.Resource;

/**
 * @Author walket
 * @create 2021/4/8 9:56 上午
 */
@Getter
public class H2SQLScriptResourceLoader extends DefaultResourceLoader {

  private List<String> tables;

  public H2SQLScriptResourceLoader() {
    this.tables = new ArrayList<>();
  }

  public H2SQLScriptResourceLoader(List<String> tables) {
    this.tables = tables;
  }

  @Override
  public Resource getResource(String location) {
    return new FileSystemResource(location) {
      @Override
      public InputStream getInputStream() throws IOException {
        //H2不支持 COLLATE 关键字以及脚本中不支持双银行"，需要替换
        String newScript = Files.readAllLines(getFile().toPath())
            .stream()
            .flatMap((Function<String, Stream<String>>) s ->
                Stream.of(s.replace("\"", "'").replace("COLLATE=utf8_bin", "").replace("ENGINE=InnoDB DEFAULT CHARSET=utf8","")
            ))
            .collect(Collectors.joining());
        StringBuilder noConstraintScript = new StringBuilder();
        //去除索引已经，不支持的约束,H2不支持索引重名，所以这里内存加载的时候取掉
        MySqlStatementParser parser = new MySqlStatementParser(newScript);
        for (SQLStatement stmt : parser.parseStatementList()) {
          if (stmt instanceof MySqlCreateTableStatement) {
            MySqlCreateTableStatement statement = (MySqlCreateTableStatement) stmt;
            tables.add(statement.getName().getSimpleName());
            List<SQLTableElement> newTableElement = statement.getTableElementList().stream()
                .filter(sqlTableElement ->
                    (sqlTableElement instanceof MySqlPrimaryKey)
                        ||
                        !(sqlTableElement instanceof MySqlTableIndex
                            || sqlTableElement instanceof MySqlKey))
                .collect(Collectors.toList());
            List<SQLTableElement> tableElementList = statement.getTableElementList();
            tableElementList.clear();
            tableElementList.addAll(newTableElement);
          } else if (stmt instanceof SQLAlterTableStatement) {
            SQLAlterTableStatement sqlAlterTableStatement = (SQLAlterTableStatement) stmt;
            List<SQLAlterTableItem> items = sqlAlterTableStatement.getItems();
            List<SQLAlterTableItem> newItems = items.stream().filter(
                sqlAlterTableItem -> !(sqlAlterTableItem instanceof SQLAlterTableDropIndex) &&
                    !(sqlAlterTableItem instanceof SQLAlterTableAddIndex) &&
                    !(sqlAlterTableItem instanceof SQLAlterTableRenameIndex)).collect(Collectors.toList());
            items.clear();
            if (newItems.isEmpty()) {
              continue;
            }
            newItems.forEach(sqlAlterTableStatement::addItem);
            System.out.println(sqlAlterTableStatement.toLowerCaseString());

          }
          String s = stmt.toLowerCaseString();
          noConstraintScript.append(s).append("\t\n");
        }
        return new ByteArrayInputStream(noConstraintScript.toString().getBytes());
      }
    };
  }
}

